import os
import time
import tkinter as tk
from tkinter import filedialog
from pathlib import Path
import threading
from tkinter import ttk  # 导入 ttk 模块用于进度条
import re
import cv2
import numpy as np

global gpu_flag
gpu_flag = 0

def central_win(win):
    win.resizable(0,0)                      # 不可缩放
    screenwidth = win.winfo_screenwidth()	# 获取屏幕分辨率宽
    screenheight = win.winfo_screenheight()	# 获取屏幕分辨率高
    win.update()	# 更新窗口
    width = win.winfo_width()	# 重新赋值
    height = win.winfo_height()
    size = '+%d+%d' % ((screenwidth - width)/2, (screenheight - height)/2)
    # 重新赋值大小 大小为屏幕大小/2
    win.geometry(size) 	# 以新大小定义窗口


def create_output_structure():
    """创建基础输出目录结构"""
    Path("output").mkdir(exist_ok=True)
    Path("videos").mkdir(exist_ok=True)
    Path("trans").mkdir(exist_ok=True)

def empty_directory(directory):
    """清空指定目录下的所有文件和子目录"""
    for root, dirs, files in os.walk(directory, topdown=False):
        for file in files:
            file_path = os.path.join(root, file)
            os.remove(file_path)
        for dir in dirs:
            dir_path = os.path.join(root, dir)
            os.rmdir(dir_path)

def extract_jpeg_frames(input_video, output_dir="output", update_progress=None):
    global stop_flag
    stop_flag = 0
    
    # 创建视频专属子目录
    video_name = Path(input_video).stem
    output_subdir = os.path.join(output_dir, video_name)
    if os.path.exists(output_subdir):
        # 若目录存在则清空
        empty_directory(output_subdir)
    else:
        # 若目录不存在则创建
        os.makedirs(output_subdir, exist_ok=True)

    print(f"导出目录绝对路径：{Path(output_subdir).resolve()}")

    # 二进制读取和解析（保持原有逻辑）
    with open(input_video, 'rb') as f:
        data = f.read()
    
    jpeg_start = b'\xff\xd8'
    jpeg_end = b'\xff\xd9'
    frame_count = 0
    i = 0
    
    size = len(data)
    if update_progress:
        update_progress(0, size)  # 初始化总大小
    print(size)
    while i < size:
        start = data.find(jpeg_start, i)
        if start == -1: break
        end = data.find(jpeg_end, start)
        if end == -1: break
        
        frame = data[start:end+2]
        frame_file = os.path.join(output_subdir, f"frame_{frame_count:06d}.jpg")
        with open(frame_file, 'wb') as f:
            f.write(frame)
            if update_progress:
                update_progress(i, size)  # 更新进度
        
        frame_count += 1
        i = end + 2
        if stop_flag:
            break
    
    return frame_count,Path(output_subdir).resolve()

def find_max_frame(directory):
    """查找目录中最大的frame序号"""
    max_num = -1
    pattern = re.compile(r'frame_(\d+)\.jpg$')
    
    for filename in os.listdir(directory):
        match = pattern.match(filename)
        if match:
            current_num = int(match.group(1))
            if current_num > max_num:
                max_num = current_num
    
    if max_num == -1:
        return 0
    return max_num + 1  # +1 因为序号从0开始

def make_jpeg_video(input_dir,output_file,start_frame=0,end_frame=0,fps=60,w=1920,h=1080,update_progress=None,code="MP4"):  
    global stop_flag
    stop_flag = 0

    global gpu_flag
    
    # 获取排序后的文件列表
    frame_files = sorted(
        [f for f in os.listdir(input_dir) if f.startswith("frame_") and f.endswith(".jpg")],
        key=lambda x: int(re.search(r'frame_(\d+)', x).group(1))
    )
    
    # 创建视频写入对象
    if code == "MP4":
        output_file = output_file + ".mp4"
        fourcc = cv2.VideoWriter_fourcc(*'avc1')
    elif code == "AVI":
        output_file = output_file + ".avi"
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
    else:
        return '',0,0
    print(output_file)

    video_writer = cv2.VideoWriter(output_file, fourcc, fps, (w, h))
    
    # 视频总帧数
    total_frames  = len(frame_files)

    if start_frame < 0 or start_frame >= total_frames:
        start_frame = 0
    if end_frame <= 0 or end_frame >= total_frames:
        end_frame = total_frames - 1

    frame_files_list = []
    frame_files_list = frame_files[start_frame : end_frame+1]

    total_frames  = len(frame_files_list)
    if update_progress:
        update_progress(0, total_frames )  # 初始化总大小
    print(total_frames )

    # 批量写入优化
    for i, filename in enumerate(frame_files_list):
        img_path = os.path.join(input_dir, filename)
        frame = cv2.imread(img_path, cv2.IMREAD_IGNORE_ORIENTATION | cv2.IMREAD_COLOR)
        if frame is not None:
            # 调整帧尺寸（如果需要）
            if frame.shape[:2] != (h, w):
                if gpu_flag:                    
                    frame = cv2.resize(cv2.UMat(frame), (w, h))
                else:
                    frame = cv2.resize(frame, (w, h))
            video_writer.write(frame)        
            
        if update_progress: 
            update_progress(i+1, total_frames )
        
        if stop_flag:
            break
    
    video_writer.release()

    return output_file,start_frame,end_frame

def trans_jpeg_video(input_file, output_file, fps=60,w=1920,h=1080,update_progress=None,code="MP4"):
    global stop_flag
    stop_flag = 0
    
    with open(input_file, 'rb') as f:
        data = f.read()

    # 创建视频写入对象
    if code == "MP4":
        output_file = output_file + ".mp4"
        fourcc = cv2.VideoWriter_fourcc(*'avc1')
    elif code == "AVI":
        output_file = output_file + ".avi"
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
    else:
        return '',0
    print(output_file)

    video_writer = cv2.VideoWriter(output_file, fourcc, fps, (w, h))

    i = 0
    frames = 0
    size = len(data)
    if update_progress:
        update_progress(0, size)  # 初始化总大小
    print(size)
    while i < size and not stop_flag:
        start = data.find(b'\xff\xd8', i)
        if start == -1: break
        end = data.find(b'\xff\xd9', start)
        if end == -1: break
        
        # 直接解码帧数据并调整尺寸
        frame = cv2.imdecode(np.frombuffer(data[start:end+2], dtype=np.uint8), cv2.IMREAD_COLOR)
        if frame is not None:
            # 统一调整帧尺寸
            if frame.shape[:2] != (h, w):
                if gpu_flag:                    
                    frame = cv2.resize(cv2.UMat(frame), (w, h))
                else:
                    frame = cv2.resize(frame, (w, h))
            video_writer.write(frame)
            frames += 1
        
        if update_progress:
            update_progress(i, size)

        i = end + 2

        if stop_flag:
            break
        
    video_writer.release()
    return output_file,frames

class VideoExtractorApp:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("MJPG 帧提取器")
        self.root.geometry("500x420")
        self.root.resizable(False, False)
        central_win(self.root)

        # 创建主容器
        self.frame = tk.Frame(self.root, padx=20, pady=20)
        self.frame.pack()
        
        # 修改按钮布局
        self.btn_frame = tk.Frame(self.frame)
        self.btn_frame.pack(pady=10)
        
        # 选择文件按钮（仅选择不导出）
        self.btn_select = tk.Button(
            self.btn_frame, 
            text="选择视频文件", 
            command=self.select_file,
            width=15
        )
        self.btn_select.pack(side=tk.LEFT, padx=5)

        self.selected_video = ""

        # 新增导出按钮
        self.btn_export = tk.Button(
            self.btn_frame,
            text="导出图片帧",
            command=self.start_export_thread,
            width=15,
            state=tk.DISABLED  # 初始禁用
        )
        self.btn_export.pack(side=tk.LEFT, padx=5)
        
        self.btn_convert = tk.Button(
            self.btn_frame,
            text="转换原视频",
            command=self.start_trans_thread,
            width=15,
            state=tk.DISABLED
        )
        self.btn_convert.pack(side=tk.LEFT, padx=5)

        self.status_label = tk.Label(
            self.frame, 
            text="准备就绪\n选择文件或目录以进行提取、转换或制作",
            fg="gray"
        )
        self.status_label.pack()

        # 添加总帧数标签
        self.total_size_label = tk.Label(
            self.frame, 
            text="总大小: 0",
            fg="black"
        )
        self.total_size_label.pack()

        # 添加进度条
        self.progress_bar = ttk.Progressbar(
            self.frame, 
            orient="horizontal",
            length=300,
            mode="determinate"
        )
        self.progress_bar.pack(pady=10)

        # 新增目录选择组件
        self.dir_frame = tk.Frame(self.frame)
        self.dir_frame.pack(pady=5)
        
        self.dir_entry = tk.Entry(self.dir_frame, width=30)
        self.dir_entry.pack(side=tk.LEFT, padx=5)
        
        self.btn_browse = tk.Button(
            self.dir_frame,
            text="浏览文件夹",
            command=self.select_directory,
            width=10
        )
        self.btn_browse.pack(side=tk.LEFT)

                # 在制作视频按钮上方添加参数选择
        self.param_frame = tk.Frame(self.frame)
        self.param_frame.pack(pady=5)
        
        # 帧率选择
        self.fps_label = tk.Label(self.param_frame, text="帧率:")
        self.fps_label.grid(row=0, column=0, padx=2)
        self.fps_combo = ttk.Combobox(self.param_frame, values=[60,45,30], width=6)
        self.fps_combo.set(60)
        self.fps_combo.grid(row=0, column=1, padx=2)
        
        # 分辨率选择
        self.res_label = tk.Label(self.param_frame, text="分辨率:")
        self.res_label.grid(row=0, column=2, padx=2)
        self.res_combo = ttk.Combobox(self.param_frame, 
                                    values=["1920x1080", "1280x720"], 
                                    width=10)
        self.res_combo.set("1920x1080")
        self.res_combo.grid(row=0, column=3, padx=2)
        
        # 格式选择
        self.format_label = tk.Label(self.param_frame, text="格式:")
        self.format_label.grid(row=0, column=4, padx=2)
        self.format_combo = ttk.Combobox(self.param_frame, 
                                       values=["MP4", "AVI"], 
                                       width=6)
        self.format_combo.set("MP4")
        self.format_combo.grid(row=0, column=5, padx=2)

        # 在制作视频按钮前添加帧范围输入
        self.range_frame = tk.Frame(self.frame)
        self.range_frame.pack(pady=5)
        
        self.start_label = tk.Label(self.range_frame, text="起始帧:")
        self.start_label.pack(side=tk.LEFT, padx=2)
        self.start_entry = tk.Entry(self.range_frame, width=6)
        self.start_entry.pack(side=tk.LEFT, padx=2)
        self.start_entry.insert(0, "0")
        
        self.end_label = tk.Label(self.range_frame, text="停止帧:")
        self.end_label.pack(side=tk.LEFT, padx=2)
        self.end_entry = tk.Entry(self.range_frame, width=6)
        self.end_entry.pack(side=tk.LEFT, padx=2)
        self.end_entry.insert(0, "0")

        # 原制作视频按钮位置保持不变
        self.btn_make_video = tk.Button(
            self.range_frame,  # 改为放入range_frame
            text="制作视频",
            command=self.start_make_video_thread,
            width=20
        )
        self.btn_make_video.pack(side=tk.LEFT, padx=10)

        self.stop_btn = tk.Button(
            self.frame,
            text="停止任务",
            command=self.stop_thread,
            width=20
        )
        self.stop_btn.pack(pady=10)

        self.time_label = tk.Label(
            self.frame, 
            text="时间记录",
            fg="gray"
        )
        self.time_label.pack()

        self.video_fames = 0
        self.selected_dir = ''

        self.current_progress = 0
        self.total_progress = 0
        
        # 添加定时器
        self.root.after(1000, self.update_progress_timer)

        self.ocl_available = cv2.ocl.haveOpenCL()
        self.ocl_device = cv2.ocl.Device_getDefault()

        # 在GPU状态显示前添加开关控件
        self.gpu_toggle_frame = tk.Frame(self.frame)
        self.gpu_toggle_frame.pack(pady=5)

        # GPU加速开关
        self.gpu_toggle = tk.Checkbutton(
            self.gpu_toggle_frame,
            text="启用GPU加速",
            command=self.toggle_gpu_acceleration,
            state=tk.DISABLED  # 初始禁用
        )
        self.gpu_toggle.pack(side=tk.LEFT)

        # 原GPU状态标签
        self.gpu_status = tk.Label(
            self.gpu_toggle_frame,
            text=self.check_gpu_acceleration(),
            fg="blue"
        )
        self.gpu_status.pack(side=tk.LEFT, padx=5)

        # 初始化输出目录
        create_output_structure()

    def toggle_gpu_acceleration(self):
        """切换GPU加速状态"""
        global gpu_flag
        gpu_flag = not gpu_flag
        status_text = f"GPU加速: {self.ocl_device.name()} {'启用' if gpu_flag else '禁用'}"
        self.gpu_status.config(text=status_text)

    def check_gpu_acceleration(self):
        """显示硬件加速状态"""
        global gpu_flag
        if self.ocl_available:
            # 仅在检测到GPU时启用开关
            self.gpu_toggle.config(state=tk.NORMAL)
            return f"GPU加速: {self.ocl_device.name()} {'禁用'}"
        # 无GPU时强制关闭加速
        gpu_flag = 0
        self.gpu_toggle.deselect()
        return "GPU加速: 不可用"
    
    def stop_thread(self):
        global stop_flag
        stop_flag = 1
        self.ctrl_btn(1)

    def ctrl_btn(self,i):
        if i :
            self.gpu_toggle.config(state=tk.NORMAL)
            self.btn_make_video.config(state=tk.NORMAL)
            self.btn_browse.config(state=tk.NORMAL)
            self.btn_select.config(state=tk.NORMAL)
            self.btn_export.config(state=tk.NORMAL)
            self.btn_convert.config(state=tk.NORMAL)
        else:
            self.gpu_toggle.config(state=tk.DISABLED)
            self.btn_make_video.config(state=tk.DISABLED)
            self.btn_browse.config(state=tk.DISABLED)
            self.btn_select.config(state=tk.DISABLED)
            self.btn_export.config(state=tk.DISABLED)
            self.btn_convert.config(state=tk.DISABLED)
    
    def trans_video_thread(self, output_path, fps, w, h, code):
        t=time.time()
        """视频转换线程"""
        self.status_label.config(text=f"正在读取... \n{os.path.basename(self.selected_video)}", fg="blue")
        try:
            result = trans_jpeg_video(self.selected_video, output_path, fps=fps, w=w, h=h, code=code,
                                    update_progress=self.update_progress)
            self.status_label.config(text=f"转换完成: {result[0]}\n使用帧数: {result[1]}", fg="green")
        except Exception as e:
            self.status_label.config(text=f"转换失败: \n{str(e)}", fg="red")
        finally:
            self.ctrl_btn(1)

        self.selected_video = ''
        self.time_label.config(text=f"上一次操作时间 {(time.time()-t)}", fg="blue")

    def start_trans_thread(self):
        """启动转换线程"""
        try:
            fps = int(self.fps_combo.get())
        except:
            fps = 60
            self.fps_combo.set("60")
        
        try:
            res = self.res_combo.get().split('x')
            w, h = int(res[0]), int(res[1])
        except:
            w, h = 1920, 1080
            self.res_combo.set("1920x1080")
            
        try:
            code = self.format_combo.get()
        except:
            code = 'MP4'
            self.format_combo.set("MP4")

        if self.selected_video:
            self.ctrl_btn(0)
            output_path = os.path.join("trans", f"{Path(self.selected_video).stem}_converted")
            thread = threading.Thread(target=self.trans_video_thread,args=(output_path, fps, w, h, code))
            thread.start()
        else:
            self.status_label.config(text="请先选择视频文件\n ", fg="red")

    def start_make_video_thread(self):
        """启动制作视频线程"""
        try:
            fps = int(self.fps_combo.get())
        except:
            fps = 60
            self.fps_combo.set("60")
        
        try:
            res = self.res_combo.get().split('x')
            w, h = int(res[0]), int(res[1])
        except:
            w, h = 1920, 1080
            self.res_combo.set("1920x1080")
            
        try:
            code = self.format_combo.get()
        except:
            code = 'MP4'
            self.format_combo.set("MP4")

        try:
            start_frame = int(self.start_entry.get())
        except:
            start_frame = 0
        
        try:
            end_frame = int(self.end_entry.get())
        except:
            end_frame = 0

        if end_frame < 0:
            end_frame = 0
        if start_frame < 0:
            start_frame = 0

        if start_frame > end_frame:
            start_frame, end_frame = end_frame, start_frame  # 交换顺序

        self.start_entry.delete(0, tk.END)
        self.start_entry.insert(0, str(start_frame))
        self.end_entry.delete(0, tk.END)
        self.end_entry.insert(0, str(end_frame))

        self.selected_dir = self.dir_entry.get()
        if not self.selected_dir:
            self.status_label.config(text="请先选择正确目录\n ", fg="red")
            self.video_fames = 0
            self.selected_dir = ''
            self.ctrl_btn(1)
            return

        self.video_fames = find_max_frame(self.selected_dir)
        self.status_label.config(
                text=f"已选择目录: {os.path.basename(self.selected_dir)} \n帧数量: {self.video_fames}",
                fg="blue"
            )
        
        if not os.path.exists(self.selected_dir) or self.video_fames == 0:
            self.status_label.config(text="请先选择正确目录\n ", fg="red")
            self.selected_dir = ''
            self.video_fames = 0
            self.ctrl_btn(1)
            return
        
        self.ctrl_btn(0)
        thread = threading.Thread(target=self.make_video,args=(code, fps, w, h,start_frame,end_frame))
        thread.start()

    def make_video(self, code_str, fps, w, h,start_frame,end_frame):
        t=time.time()
        """视频生成逻辑"""
        try:
            # 这里添加实际的视频生成代码
            # 使用 self.selected_dir 和 self.video_fames
            output_path = os.path.join("videos", f"{os.path.basename(self.selected_dir)}")
            make_result = make_jpeg_video(self.selected_dir,output_path,
                                       start_frame=start_frame,end_frame=end_frame,
                                       fps=fps, w=w, h=h,
                                       update_progress=self.update_progress,code=code_str)
            # 示例结果
            if make_result[0]:
                result = f"已生成视频: {make_result[0]}\n使用帧数: {self.video_fames}"
            else:
                result = f"视频格式错误: {output_path}.{code_str}\n "
            
            self.start_entry.delete(0, tk.END)
            self.start_entry.insert(0, str(make_result[1]))
            self.end_entry.delete(0, tk.END)
            self.end_entry.insert(0, str(make_result[2]))

            self.status_label.config(text=result, fg="green")
        except Exception as e:
            self.status_label.config(text=f"生成失败: \n{str(e)}", fg="red")
        finally:
            self.ctrl_btn(1)
        
        self.time_label.config(text=f"上一次操作时间 {(time.time()-t)}", fg="blue")

    def select_file(self):
        """仅选择文件不导出"""
        self.ctrl_btn(0)
        file_path = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[("AVI 视频", "*.avi"),("MJPEG 视频", "*.mjpg;*.mjpeg"), ("所有文件", "*.*")]
        )
        
        if file_path:
            self.selected_video = file_path  # 存储选择的文件路径
            self.btn_export.config(state=tk.NORMAL)  # 启用导出按钮
            self.status_label.config(
                text=f"已选择文件: \n{os.path.basename(file_path)}",
                fg="blue"
            )
        self.ctrl_btn(1)
    
    def start_export_thread(self):
        """启动导出线程"""
        self.status_label.config(text=f"正在读取... \n{os.path.basename(self.selected_video)}", fg="blue")
        if self.selected_video:
            self.ctrl_btn(0)
            thread = threading.Thread(target=self.extract_frames_thread, args=(self.selected_video,))
            thread.start()
        else:
            self.status_label.config(text="请先选择视频文件\n ", fg="red")
        
        self.selected_video = ''

    def select_directory(self):
        self.video_fames = 0        
        """打开文件夹选择对话框"""
        # 设置默认目录为当前文件所在目录下的 output 文件夹
        default_dir = os.path.abspath("output")  # 获取output目录的绝对路径
        self.dir_entry.delete(0, tk.END)
        self.selected_dir = ''
        self.selected_dir = filedialog.askdirectory(
            title="选择输出目录",
            initialdir=default_dir  # 添加初始目录参数
        )
        if self.selected_dir:
            self.video_fames = find_max_frame(self.selected_dir)
            self.dir_entry.delete(0, tk.END)
            self.dir_entry.insert(0, self.selected_dir)
            self.status_label.config(
                text=f"已选择目录: {os.path.basename(self.selected_dir)} \n帧数量: {self.video_fames}",
                fg="blue"
            )

    def extract_frames_thread(self, file_path):
        t=time.time()
        try:
            result = extract_jpeg_frames(file_path, update_progress=self.update_progress)
            self.selected_dir = result[1]
            self.video_fames = result[0]
            if self.selected_dir and self.video_fames:
                self.dir_entry.delete(0, tk.END)
                self.dir_entry.insert(0, self.selected_dir)
                self.status_label.config(
                    text=f"成功导出 {self.video_fames} 帧到目录: \n{os.path.basename(self.selected_dir)}"
                    ,
                    fg="green"
                )
            else:
                self.dir_entry.delete(0, tk.END)
                self.status_label.config(
                    text=f"失败 {self.video_fames} 帧到目录: \n{os.path.basename(self.selected_dir)}"
                    ,
                    fg="red"
                )
        except Exception as e:
            self.status_label.config(
                text=f"提取失败: \n{str(e)}",
                fg="red"
            )
        finally:
            self.ctrl_btn(1)
        self.time_label.config(text=f"上一次操作时间 {(time.time()-t)}", fg="blue")

    def update_progress_timer(self):
        """定时器回调函数"""
        self.total_size_label.config(text=f"总大小: {self.total_progress}")
        progress = (self.current_progress / self.total_progress) * 100 if self.total_progress > 0 else 0
        self.progress_bar['value'] = progress
        
        # 重新启动定时器
        self.root.after(1000, self.update_progress_timer)

    def update_progress(self, current, total):
        """更新进度数据（不再直接操作UI）"""
        self.current_progress = current
        self.total_progress = total

if __name__ == "__main__":
    global stop_flag
    app = VideoExtractorApp()
    app.root.mainloop()
    stop_flag = 1
